package ds.test;

import org.dbunit.database.DatabaseConfig;
import org.unitils.dbunit.DbUnitModule;
import org.unitils.dbunit.annotation.DataSet;
import org.unitils.dbunit.annotation.ExpectedDataSet;
import org.unitils.dbunit.datasetfactory.DataSetFactory;
import org.unitils.dbunit.util.DbUnitDatabaseConnection;
import org.unitils.dbunit.util.MultiSchemaDataSet;

import java.lang.reflect.Method;
import java.util.Map;

public class DBUnitModule extends DbUnitModule {

    @Override
    public DbUnitDatabaseConnection getDbUnitDatabaseConnection(final String schemaName) {
        DbUnitDatabaseConnection dbConnection = super.getDbUnitDatabaseConnection(schemaName);
        dbConnection.getConfig().setProperty(DatabaseConfig.PROPERTY_TABLE_TYPE, new String[]{"VIEW", "TABLE"});
        return dbConnection;
    }

    @Override
    public MultiSchemaDataSet getDataSet(DataSet dataSet, Method testMethod, Class<?> testClass) {
        databaseName = dataSet.databaseName();
        // Create configured factory for data sets
        String[] dataSetFileNames = dataSet.value();
        convertExcelFactory(dataSetFileNames);
        DataSetFactory dataSetFactory = getDataSetFactory(DataSet.class, testMethod, testClass);

        // Get the dataset file name
        if (dataSetFileNames.length == 0) {
            // empty means, use default file name, which is the name of the class + extension
            dataSetFileNames = new String[]{
                    getCorrectFileName(testClass, testMethod, dataSetFactory.getDataSetFileExtension())
            };
        }
        return getDataSet(testClass, dataSetFileNames, dataSetFactory);
    }

    private void convertExcelFactory(String[] dataSetFileNames) {
        boolean isExcelData = false;
        for (String dataSetFileName : dataSetFileNames) {
            if (dataSetFileName.lastIndexOf("xlsx") > -1 || dataSetFileName.lastIndexOf("xls") > -1) {
                isExcelData = true;
                break;
            }
        }

        if (isExcelData) {
            Map<String, String> dataSetPropertyValues = defaultAnnotationPropertyValues.get(DataSet.class);
            dataSetPropertyValues.put("factory", "ds.test.MultiSchemaXlsDataSetFactory");
            Map<String, String> expectedDataSetPropertyValues = defaultAnnotationPropertyValues.get(ExpectedDataSet.class);
            expectedDataSetPropertyValues.put("factory", "ds.test.MultiSchemaXlsDataSetFactory");
        }
    }
}
